from flask import Flask, request, jsonify
from PIL import Image
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import io

app = Flask(__name__)


class MNISTPredictor:
    def __init__(self, model_path):
        self.model = self.load_model(model_path)
        self.transform = transforms.Compose([
            transforms.Resize((28, 28)),
            transforms.ToTensor(),
            transforms.Lambda(lambda x: 1 - x if x.mean() > 0.5 else x)
        ])

    def load_model(self, path):
        # 必须与训练代码中的CNN类完全一致
        class CNN(nn.Module):
            def __init__(self):
                super(CNN, self).__init__()
                self.conv1 = nn.Sequential(
                    nn.Conv2d(1, 32, 3, 1, 1),
                    nn.ReLU(),
                    nn.MaxPool2d(2, 2)
                )
                self.conv2 = nn.Sequential(
                    nn.Conv2d(32, 64, 3, 1, 1),
                    nn.ReLU(),
                    nn.MaxPool2d(2, 2)
                )
                self.fc1 = nn.Sequential(
                    nn.Linear(64 * 7 * 7, 256),
                    nn.ReLU(),
                    nn.Dropout(p=0.5)
                )
                self.fc2 = nn.Linear(256, 10)

            def forward(self, x):
                x = self.conv1(x)
                x = self.conv2(x)
                x = x.view(x.size(0), -1)
                x = self.fc1(x)
                x = self.fc2(x)
                return x

        model = CNN()
        model.load_state_dict(torch.load(path, map_location='cpu', weights_only=False))
        model.eval()
        return model

    def predict(self, image_bytes):
        img = Image.open(io.BytesIO(image_bytes)).convert('L')
        tensor = self.transform(img).unsqueeze(0)
        with torch.no_grad():
            outputs = self.model(tensor)
            probs = torch.softmax(outputs, dim=1)[0]
            pred = torch.argmax(probs).item()
        return {
            'prediction': int(pred),
            'probabilities': [round(p.item(), 4) for p in probs]
        }


class Result:
    def __init__(self, code=200, message='success', data=None):
        self.code = code
        self.message = message
        self.data = data

    def to_dict(self):
        return {
            'code': self.code,
            'message': self.message,
            'data': self.data
        }


predictor = MNISTPredictor('../model_save/mnist_cnn_net02.pth')


@app.route('/predict', methods=['POST'])
def predict():
    if 'image' not in request.files:
        return jsonify(Result(400, 'No image uploaded').to_dict())

    try:
        file = request.files['image']
        if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            return jsonify(Result(400, 'Invalid image format').to_dict())

        result = predictor.predict(file.read())
        return jsonify(Result(data=result).to_dict())
    except Exception as e:
        return jsonify(Result(500, str(e)).to_dict())


if __name__ == '__main__':
    app.run(port=5000)